package com.aiwiown.snackmq.cluster.timer;

import com.aiwiown.snackmq.cluster.raft.entity.SnackMQRaftLogEntry;
import com.aiwiown.snackmq.cluster.raft.manager.SnackMQRaftNodeManager;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.protocol.OperationType;
import com.aiwiown.snackmq.storage.service.PartitionedStorageService;
import com.aiwiown.snackmq.storage.timer.TimerMessageStore;
import lombok.extern.slf4j.Slf4j;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 后台服务，定期扫描时间轮，处理到期的延迟消息。
 */
@Slf4j
public class TimerScannerService {

    private final ScheduledExecutorService scannerScheduler;
    private final TimerMessageStore timerMessageStore;
    private final PartitionedStorageService storageService;
    private final SnackMQRaftNodeManager dataRaftNodeManager;
    private final int wheelSize;

    public TimerScannerService(TimerMessageStore timerMessageStore,
                               PartitionedStorageService storageService,
                               SnackMQRaftNodeManager dataRaftNodeManager,
                               int wheelSize) {
        this.timerMessageStore = timerMessageStore;
        this.storageService = storageService;
        this.dataRaftNodeManager = dataRaftNodeManager;
        this.wheelSize = wheelSize;
        this.scannerScheduler = Executors.newSingleThreadScheduledExecutor(r -> new Thread(r, "TimerScannerThread"));
    }

    public void start() {
        // 每秒执行一次，模拟时间轮指针前进
        scannerScheduler.scheduleAtFixedRate(this::scanAndDeliver, 1, 1, TimeUnit.SECONDS);
        log.info("TimerScannerService started.");
    }

    public void shutdown() {
        scannerScheduler.shutdown();
        try {
            if (!scannerScheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                scannerScheduler.shutdownNow();
            }
        } catch (InterruptedException e) {
            scannerScheduler.shutdownNow();
            Thread.currentThread().interrupt();
        }
        log.info("TimerScannerService shutdown.");
    }

    private void scanAndDeliver() {
        // 确保只在 Leader 节点上执行扫描和投递
        if (!dataRaftNodeManager.isLeader()) {
            return;
        }
        try {
            long now = System.currentTimeMillis();
            int slotToScan = (int) ((now / 1000) % wheelSize);

            // 1. 使用新的、更简单的方法从存储中获取所有候选条目
            List<TimerMessageStore.TimerLogEntry> entriesInSlot = timerMessageStore.fetchAndClearSlotEntries(slotToScan);
            if (entriesInSlot.isEmpty()) {
                return;
            }
            log.debug("Scanning slot {}. Found {} candidate entries.", slotToScan, entriesInSlot.size());

            for (TimerMessageStore.TimerLogEntry entry : entriesInSlot) {
                long deliveryTimestamp = entry.getDeliveryTimestamp();
                long delay = deliveryTimestamp - now;

                if (delay <= 0) {
                    // 消息已到期，立即重新提交
                    resubmitMessage(entry);
                } else {
                    // 消息尚未到期，安排一个精确的延迟任务
                    log.debug("Scheduling message from offset {} for precise delivery in {} ms.", entry.getCommitLogOffset(), delay);
                    scannerScheduler.schedule(() -> resubmitMessage(entry), delay, TimeUnit.MILLISECONDS);
                }
            }
        } catch (Exception e) {
            log.error("Error during timer scan and delivery.", e);
        }
    }

    /**
     * 【新增辅助方法】: 封装了恢复和重新提交消息的逻辑，以避免代码重复。
     *
     * @param entry 延迟消息的日志条目
     */
    private void resubmitMessage(TimerMessageStore.TimerLogEntry entry) {
        try {
            // 1. 从 Raft 日志中恢复原始消息
            SnackMQRaftLogEntry logEntry = dataRaftNodeManager.readLogEntry(entry.getCommitLogOffset());

            if (logEntry == null) {
                log.error("CRITICAL: Failed to recover delayed message from Raft log at offset {}. The message may be lost.",
                        entry.getCommitLogOffset());
                return;
            }

            Message originalBatch = logEntry.getData(Message.class);
            if (originalBatch == null || !originalBatch.isBatch() || originalBatch.getMessages().isEmpty()) {
                log.error("CRITICAL: Recovered Raft log entry at offset {} has null or invalid data. The message may be lost.",
                        entry.getCommitLogOffset());
                return;
            }

            Message originalMessage = originalBatch.getMessages().get(0);

            // 2. 准备重新投递的消息：清除延迟属性
            Message readyToDeliver = originalMessage.toBuilder()
                    .property("delayTimeInSeconds", "0") // 覆盖延迟属性
                    .build();

            // 3. 再次包装成批处理消息
            Message batchForResubmit = Message.builder()
                    .type(MessageType.PRODUCE_REQUEST)
                    .messages(Collections.singletonList(readyToDeliver))
                    .build();

            // 4. 重新提交到 Raft
            SnackMQRaftLogEntry resubmitLogEntry = new SnackMQRaftLogEntry(OperationType.PRODUCE, batchForResubmit);

            dataRaftNodeManager.apply(resubmitLogEntry, status -> {
                if (status.isOk()) {
                    log.info("Successfully re-submitted delayed message from offset {} to topic {}.",
                            entry.getCommitLogOffset(), readyToDeliver.getTopic());
                } else {
                    log.error("Failed to re-submit delayed message from offset {}. Status: {}",
                            entry.getCommitLogOffset(), status);
                }
            });
        } catch (Exception e) {
            log.error("Failed to resubmit message from entry: {}", entry, e);
        }
    }
}